import wordcloud
import jieba
import imageio
from utils.commonUtils import get_all_file_in_dir


def wordcloud_trans(input_path, output_path, word_path, font='华文 彩云'):
    '''
    图片批量转词云
    :param input_path: 输入图片的文件夹路径
    :param output_path: 输出图片的文件夹路径
    :param word_path: 词云文件路径
    :return:
    '''
    img_path_list = get_all_file_in_dir(input_path)
    for i, img_path in enumerate(img_path_list):
        drow_mask_wordColud(input_img_path=img_path, output_img_path=img_path.replace(input_path, output_path), word_path=word_path, font=font)
        print('word_cloud_trans: {0}%'.format(round((i + 1) / len(img_path_list) * 100, 2)))


def drow_mask_wordColud(input_img_path, output_img_path, word_path, font):
    '''
    单个图片转词云
    :param input_path: 输入图片路径
    :param output_path: 输出图片路径
    :param word_path: 词云文件路径
    :return:
    '''

    # 字体字典
    font_root = "C:/Windows/Fonts/"
    font_dict = {
        "幼圆 常规": "SIMYOU.TTF",
        "华文 彩云": "华文彩云.ttf",
        "微软雅黑 常规": "msyh.ttc",
        "方正粗黑宋简体": "方正粗黑宋简体.ttf",
    }

    # 1、准备文本
    f = open(word_path, encoding='utf-8')
    kkx = f.read()
    kkx = jieba.cut(kkx)
    kkx = " ".join(kkx)
    # 2、生成图片的nd-array，传入图片路径
    im = imageio.imread(input_img_path)
    # 3、获取一个图形颜色生成器
    image_color = wordcloud.ImageColorGenerator(im)
    # 4、创建词云对象
    wc = wordcloud.WordCloud(background_color='white', mask=im, contour_color='black', max_font_size=50, min_font_size=5, max_words=2000, font_path=font_root+font_dict[font])
    # 5、根据文本生成词云
    wc.generate(kkx)
    # 根据图片颜色重绘
    rwc = wc.recolor(color_func=image_color)
    rwc.to_file(output_img_path)


if __name__ == '__main__':
    drow_mask_wordColud(input_img_path='../production_process/3_style_change_imgs/1_20.jpg', output_img_path='qmkkx.png', word_path='../text/test.txt', font='华文 彩云')